import torch
import torch_geometric.nn
from torch_geometric.nn import HeteroConv, SAGEConv, GATConv, Linear
from torch.nn.functional import sigmoid



class MyGNN(torch.nn.Module):
    def __init__(self, hidden_channels, num_layers):
        super().__init__()

        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HeteroConv({
                ('var', 'v_c', 'con'): SAGEConv((-1, -1), hidden_channels),
                ('con', 'c_v', 'var'): SAGEConv((-1, -1), hidden_channels),
                ('var', 'v_m', 'met'): SAGEConv((-1, -1), hidden_channels),
                ('met', 'm_v', 'var'): SAGEConv((-1, -1), hidden_channels),
            }, aggr='mean')
            self.convs.append(conv)

        self.lin = Linear(in_channels=hidden_channels, out_channels=1)

    def forward(self, x_dict, edge_index_dict, num_graphs):
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}
        logit = self.lin(x_dict['var'])  # shape: (num_var_node*num_graphs,1)

        logit = logit.reshape([num_graphs, -1]).mean(dim=-1)  # shape: (num_graphs,)
        return sigmoid(logit)


if __name__ == '__main__':
    mygnn = MyGNN(hidden_channels=4, num_layers=2)
    # g2 = torch_geometric.nn.PNAConv()
    # g3 = torch_geometric.nn.CGConv()
